import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import random

plt.rcParams['font.sans-serif'] = ['SimHei']

POP_SIZE = 100  # 种群大小
GENERATION_SIZE = 100  # 迭代次数
CROSS_RATE = 0.8  # 交叉概率
MUTATION_RATE = 0.003  # 变异概率
X_BOUND = [-32, 32]  # x的范围
DNA_SIZE = 22  # DNA长度（解的二进制长度，保留了6位小数）


# 目标函数
def F(x1, x2):
    d = 2
    a, b, c = 20, 0.2, 2 * np.pi
    term1 = -a * np.exp(-b * np.sqrt((x1 ** 2 + x2 ** 2) / d))
    term2 = -np.exp((np.cos(c * x1) + np.cos(c * x2)) / d)
    result = term1 + term2 + a + np.exp(1)
    return result


# 初始化
def initialize(pop_size, dna_size):
    population = np.random.randint(0, 2, size=(pop_size, dna_size * 2))  # matrix (POP_SIZE, DNA_SIZE*2)
    return population


# 计算适应度
def get_fitness(pop):
    x1, x2 = decode(pop)
    pred = F(x1, x2)
    # 最小化问题
    return -(pred - np.max(pred)) + 1e-6


# 选择算子
def select(pop, fitness):
    idx = np.random.choice(np.arange(POP_SIZE), size=POP_SIZE, replace=True, p=(fitness / fitness.sum()))
    return pop[idx]


# 交叉算子
def crossover_mutate(pop, crossover_rate, mutate_rate):
    new_pop = []
    for father in range(pop.shape[0]):  # 遍历种群中的每一个个体，将该个体作为父亲
        child = pop[father, :]  # 孩子得到父亲的全部基因
        if np.random.rand() < crossover_rate:  # 产生子代时不是必然发生交叉，而是以一定的概率发生交叉
            mother = pop[np.random.randint(POP_SIZE), :]  # 再种群中选择另一个个体，并将该个体作为母亲
            cross_points = np.random.randint(0, DNA_SIZE * 2)  # 随机产生交叉的点
            child[cross_points:] = mother[cross_points:]  # 孩子得到位于交叉点后的母亲的基因
        child = mutate(child, mutate_rate)  # 每个后代有一定的机率发生变异
        new_pop.append(child)
    return np.array(new_pop)


# 变异算子
def mutate(child, mutate_rate):
    if np.random.rand() < mutate_rate:
        mutate_point = np.random.randint(0, DNA_SIZE * 2)
        child[mutate_point] = child[mutate_point] ^ 1
    return child


# 解码 : 二进制转十进制
def decode(pop):
    x1_pop = pop[:, 1::2]  # 偶数列
    x2_pop = pop[:, ::2]  # 奇数列
    x1 = x1_pop.dot(2 ** np.arange(DNA_SIZE)[::-1]) * (X_BOUND[1] - X_BOUND[0]) / (float(2 ** DNA_SIZE) - 1) + X_BOUND[
        0]
    x2 = x2_pop.dot(2 ** np.arange(DNA_SIZE)[::-1]) * (X_BOUND[1] - X_BOUND[0]) / (float(2 ** DNA_SIZE) - 1) + X_BOUND[
        0]
    return x1, x2


def draw_line(generation_size, z1, ylabel1, z2, ylabel2):
    plt.figure()
    plt.xlabel('迭代次数')
    plt.ylabel('种群位置')
    plt.plot(range(generation_size), z1, c='r', label=ylabel1)
    plt.plot(range(generation_size), z2, c='b', label=ylabel2)  # 最优变化曲线
    plt.legend(loc='upper right')


# 把数据保存到csv
def save_csv(idx, dna1, dna2, fit):
    # idx:迭代次数， dna1: x1, dna2: x2 , fit: fitness概率
    pass

# 最优种群分布图
def pop_distribution_optimal_graph(points_optimal):
    x1 = np.linspace(-32, 32, 50)
    x2 = np.linspace(-32, 32, 50)
    x1, x2 = np.meshgrid(x1, x2)
    z = F(x1, x2)
    fig = plt.figure()
    ax = Axes3D(fig)
    ax.plot_surface(x1, x2, z, alpha=0.4)  # 函数图
    ax.scatter(points_optimal[:, 0], points_optimal[:, 1], points_optimal[:, 2], c='r')  # 最优点

# 随机生成颜色
def randomcolor():
    colArr = ['1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F']
    color = ""
    for i in range(6):
        color += colArr[random.randint(0, 14)]
    return "#" + color

# 画出不同种群大小的最优迭代图
def draw_iter_pop(pop_size_list, pop_optimal_distribution):
    if type(pop_optimal_distribution) == list:
        pop_optimal_distribution = np.array(pop_optimal_distribution)
    # print(pop_optimal_distribution.shape)  # (5, 100, 3)
    plt.figure()
    plt.xlabel('迭代次数')
    plt.ylabel('种群位置')
    for i, num in enumerate(pop_size_list):
        plt.plot(range(GENERATION_SIZE), pop_optimal_distribution[i, :, 2], color=randomcolor(), label=f'初始种群数量:{num}')
    plt.legend()


if __name__ == '__main__':
    pop_size_list = [30, 50, 70, 100, 200]  # 不同种群大小
    pop_optimal_distribution = []           # 不同种群大小的每次迭代中的最优解 x1, x2, z
    optimal_fitness = []                    # 不同种群大小的个体的适应度概率  前25个
    for _, POP_SIZE in enumerate(pop_size_list):
        pop = initialize(POP_SIZE, DNA_SIZE)
        points_optimal = []                 # 一次种群大小的最优解集
        for i in range(GENERATION_SIZE):
            x1, x2 = decode(pop)
            # 计算个体适应度
            fitness = get_fitness(pop)
            max_idx = np.argmax(fitness)
            points_optimal.append([x1[max_idx], x2[max_idx], F(x1[max_idx], x2[max_idx])])
            points_mean.append([x1.mean(), x2.mean(), F(x1.mean(), x2.mean())])
            print('第{}次迭代：{},最大索引{}'.format(i, F(x1[max_idx], x2[max_idx]), max_idx))
            # 选择最优个体
            pop = select(pop, fitness)
            # 交叉和变异
            pop = crossover_mutate(pop, CROSS_RATE, MUTATION_RATE)

        pop_optimal_distribution.append(points_optimal)

    # 绘制
    draw_iter_pop(pop_size_list, pop_optimal_distribution)
    plt.show()
